"""
# 对使用 one-hot 编码、虚拟编码和效果编码的分类变量进行线性回归
"""

import pandas as pd
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import OneHotEncoder

df = pd.DataFrame({'City': ['SF', 'SF', 'SF',
                            'NYC', 'NYC', 'NYC',
                            'Seattle', 'Seattle', 'Seattle'],
                   'Rent': [3999, 4000, 4001, 3499, 3500, 3501, 2499, 2500, 2501]})

one_hot_df = pd.get_dummies(data=df, prefix=['City'])
dummy_df = pd.get_dummies(data=df, prefix=['City'], drop_first=True)
effect_df = dummy_df.copy()
effect_df.loc[3:5, ['City_SF', 'City_Seattle']] = -1.0

model = LinearRegression()
model.fit(one_hot_df[['City_NYC', 'City_SF', 'City_Seattle']], one_hot_df['Rent'])
print('Coef: {}\n'
      'Intercept: {}'.format(model.coef_, model.intercept_))
model.fit(dummy_df[['City_SF', 'City_Seattle']], dummy_df['Rent'])
print('Coef: {}\n'
      'Intercept: {}'.format(model.coef_, model.intercept_))
model.fit(effect_df[['City_SF', 'City_Seattle']], effect_df['Rent'])
print('Coef: {}\n'
      'Intercept: {}'.format(model.coef_, model.intercept_))

# One-hot Encoder by sklearn
encoder = OneHotEncoder(sparse=False)
res = encoder.fit_transform(df[['City']])